import os

from torch.utils.data import Dataset
import torchvision
from PIL.Image import Image
import PIL
import numpy as np


class NewMNIST(torchvision.datasets.MNIST):
    def __init__(self, root, train=True, transform=None, target_transform=None,
                 download=False):
        super().__init__(root, train, None, None, download)
        self.new_mnist_transform = transform
        self.new_mnist_target_transform = target_transform
        self.data = [
            super(NewMNIST, self).__getitem__(i)
            for i in range(super(NewMNIST, self).__len__())
        ]
        eight_imgs = list(map(
            lambda t: t[0].rotate(90, expand=0),
            filter(lambda t: t[1] == 8, self.data)
        ))
        j = 0
        for i in range(len(self.data)):
            if self.data[i][1] != 0:
                continue
            new_img = PIL.Image.fromarray(np.maximum(
                np.array(self.data[i][0]),
                np.array(eight_imgs[j])
            ))
            self.data[i] = (new_img, 0)
            j = (j + 1) % len(eight_imgs)

    @property
    def raw_folder(self):
        return os.path.join(self.root, "MNIST", 'raw')

    @property
    def processed_folder(self):
        return os.path.join(self.root, "MNIST", 'processed')

    def __getitem__(self, index):
        img, target = self.data[index]
        if self.new_mnist_transform is not None:
            img = self.new_mnist_transform(img)

        if self.new_mnist_target_transform is not None:
            target = self.new_mnist_target_transform(target)

        return img, target

    def __len__(self):
        return len(self.data)
